
from llama_index.core.indices.property_graph import SimpleLLMPathExtractor
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from llama_index.core.indices.property_graph.transformations.implicit import (
    ImplicitPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.schema_llm import (
    SchemaLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.simple_llm import (
    SimpleLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.dynamic_llm import (
    DynamicLLMPathExtractor,
)
# 初始化LLM
from llama_index.core.indices.property_graph.base import PropertyGraphIndex
from llama_index.core.indices.property_graph.retriever import PGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.base import BasePGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.custom import (
    CustomPGRetriever,
    CUSTOM_RETRIEVE_TYPE,
)
from llama_index.core.indices.property_graph.sub_retrievers.cypher_template import (
    CypherTemplateRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.llm_synonym import (
    LLMSynonymRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.text_to_cypher import (
    TextToCypherRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.vector import (
    VectorContextRetriever,
)
from llama_index.core.indices.property_graph.transformations.implicit import (
    ImplicitPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.schema_llm import (
    SchemaLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.simple_llm import (
    SimpleLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.dynamic_llm import (
    DynamicLLMPathExtractor,
)
from llama_index.core.indices.property_graph.utils import default_parse_triplets_fn

 
# 自定义提取提示词（更聚焦技术领域）
custom_prompt = (
    "分析以下技术文本，提取最多{max_paths_per_chunk}个技术相关的三元组，"
    "格式为`主体,关系,客体`，例如`LlamaIndex,是,开源框架`。避免通用词汇。\n"
    "文本：{text}\n"
    "三元组："
)
 
# 定义解析函数（处理可能的格式异常）
def parse_tech_triples(response_str):
    lines = response_str.strip().split("\n")
    print(lines)
    triples = []
    for line in lines:
        if "," in line:
            try:
                subj, pred, obj = line.split(",", 2)
                triples.append((subj.strip(), pred.strip(), obj.strip()))
            except:
                continue
    return triples
 

documents=[Document(text='''以下是技术相关的三元组（主语-谓语-宾语）示例，涵盖不同技术领域：

‌人工智能‌

Transformer模型-采用-自注意力机制
深度学习框架-支持-分布式训练
生成对抗网络（GAN）-生成-逼真图像
‌云计算‌

Kubernetes-管理-容器化应用
虚拟化技术-实现-资源隔离
Serverless架构-降低-运维复杂度''')] 

print("sssssssssssss")

# 创建提取器
kg_extractor = SimpleLLMPathExtractor(
    llm=llm,
    extract_prompt=custom_prompt,
    parse_fn=parse_tech_triples,
    max_paths_per_chunk=5,  # 每个块最多提取5条关系
    num_workers=4,          # 并行处理提高效率
)

 
# 构建带自定义提取器的索引
index = PropertyGraphIndex.from_documents(

    documents,
    llm=llm,
    kg_extractors=[kg_extractor]
)

index.property_graph_store.persist()


query_engine=index.as_query_engine()
nodes=query_engine.query("虚拟化技术")
print(nodes)